{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# MNIST\n",
    "\n",
    "MNIST is a simple computer vision dataset that consists of images of handwritten digits.\n",
    "\n",
    "Some examples:\n",
    "\n",
    "![MNIST digits sample](https://github.com/user-attachments/assets/996c11e0-47f9-4b21-8e23-3867b8942e64)\n",
    "\n",
    "It also includes labels for each image, which we use to train our example models.\n",
    "\n",
    "## The `mnist` library\n",
    "\n",
    "This package includes the following functionality:\n",
    "\n",
    "  - Download the dataset from [storage.googleapis.com/cvdf-datasets/mnist](https://storage.googleapis.com/cvdf-datasets/mnist),\n",
    "  - Create a `Dataset` object to iterate over it, use for training and evaluation.\n",
    "  - A linear and a CNN model demo.\n",
    "  - A command-line demo (in the `demo` sub-directory).\n",
    "\n",
    "This notebook serves as documentation and example for the [github.com/gomlx/gomlx/examples/mnist](https://github.com/gomlx/gomlx/examples/mnist) library, and the demo code in one piece can be seen in [.../examples/mnist/demo/](https://github.com/gomlx/gomlx/tree/main/examples/mnist/demo)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Environment Set Up\n",
    "\n",
    "Let's set up `go.mod` to use the local copy of GoMLX, so it can be developed jointly the dataset code with the model. That's often how data pre-processing and model code is developed together with experimentation.\n",
    "\n",
    "If you are not changing code, feel free to simply skip this cell. Or if you used a different directory for you projects, change it below.\n",
    "\n",
    "Notice the directory `${HOME}/Projects/gomlx` is where the GoMLX code is copied by default in [its Docker](https://hub.docker.com/repository/docker/janpfeifer/gomlx_jupyterlab/general)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\t- Added replace rule for module \"github.com/gomlx/gomlx\" to local directory \"/home/janpf/Projects/gomlx\".\n"
     ]
    }
   ],
   "source": [
    "!*rm -f go.work && go work init && go work use . \"${HOME}/Projects/gomlx\"\n",
    "%goworkfix"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Data Preparation\n",
    "\n",
    "### Downloading data files\n",
    "\n",
    "To download to the local directory, simply do the following. Notice if it's already downloaded in the given `--data` directory, it returns immediately."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "import (\n",
    "    \"github.com/gomlx/gomlx/examples/mnist\"\n",
    "    \"github.com/gomlx/gomlx/pkg/support/fsutil\"\n",
    "    \"github.com/janpfeifer/must\"\n",
    ")\n",
    "\n",
    "var flagDataDir = flag.String(\"data\", \"~/work/mnist\", \"Directory to cache downloaded and generated dataset files.\")\n",
    "\n",
    "func AssertDownloaded() {\n",
    "    *flagDataDir = must.M1(fsutil.ReplaceTildeInDir(*flagDataDir))\n",
    "    if !fsutil.MustFileExists(*flagDataDir) {\n",
    "        must.M(os.MkdirAll(*flagDataDir, 0777))\n",
    "    }\n",
    "   must.M(mnist.Download(*flagDataDir))\n",
    "}\n",
    "\n",
    "%%\n",
    "AssertDownloaded()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "total 12M\n",
      "drwxr-x--- 2 janpf janpf 4.0K Oct 11 11:33 cnn\n",
      "drwxr-x--- 2 janpf janpf 4.0K Jun  3 15:39 cnn_triplet\n",
      "drwxr-x--- 2 janpf janpf 4.0K Jun  3 15:38 linear\n",
      "drwxr-x--- 2 janpf janpf 4.0K Feb 13  2025 linear_baseline\n",
      "-rw-r--r-- 1 janpf janpf 1.6M Jun  8 08:10 t10k-images-idx3-ubyte.gz\n",
      "-rw-r--r-- 1 janpf janpf 4.5K Jun  8 08:11 t10k-labels-idx1-ubyte.gz\n",
      "-rw-r--r-- 1 janpf janpf 9.5M Jun  8 08:09 train-images-idx3-ubyte.gz\n",
      "-rw-r--r-- 1 janpf janpf  29K Jun  8 08:09 train-labels-idx1-ubyte.gz\n"
     ]
    }
   ],
   "source": [
    "!ls -lh ~/work/mnist/"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Sample some images\n",
    "The `mnist.NewDataset` creates a `data.InMemoryDataset` that can be used both for training, evaluation, or just to sample a few examples, which we do below:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<p>Samples MNIST</p>\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<table><tr>\n",
       "<td><figure style=\"padding:4px;text-align: center;\"><img width=\"56\" height=\"56\" src=\"\"><figcaption style=\"text-align: center;\">([4])</figcaption></figure></td>\n",
       "<td><figure style=\"padding:4px;text-align: center;\"><img width=\"56\" height=\"56\" src=\"\"><figcaption style=\"text-align: center;\">([2])</figcaption></figure></td>\n",
       "<td><figure style=\"padding:4px;text-align: center;\"><img width=\"56\" height=\"56\" src=\"\"><figcaption style=\"text-align: center;\">([3])</figcaption></figure></td>\n",
       "<td><figure style=\"padding:4px;text-align: center;\"><img width=\"56\" height=\"56\" src=\"\"><figcaption style=\"text-align: center;\">([0])</figcaption></figure></td>\n",
       "<td><figure style=\"padding:4px;text-align: center;\"><img width=\"56\" height=\"56\" src=\"\"><figcaption style=\"text-align: center;\">([7])</figcaption></figure></td>\n",
       "<td><figure style=\"padding:4px;text-align: center;\"><img width=\"56\" height=\"56\" src=\"\"><figcaption style=\"text-align: center;\">([7])</figcaption></figure></td>\n",
       "<td><figure style=\"padding:4px;text-align: center;\"><img width=\"56\" height=\"56\" src=\"\"><figcaption style=\"text-align: center;\">([6])</figcaption></figure></td>\n",
       "<td><figure style=\"padding:4px;text-align: center;\"><img width=\"56\" height=\"56\" src=\"\"><figcaption style=\"text-align: center;\">([3])</figcaption></figure></td>\n",
       "<td><figure style=\"padding:4px;text-align: center;\"><img width=\"56\" height=\"56\" src=\"\"><figcaption style=\"text-align: center;\">([1])</figcaption></figure></td>\n",
       "<td><figure style=\"padding:4px;text-align: center;\"><img width=\"56\" height=\"56\" src=\"\"><figcaption style=\"text-align: center;\">([1])</figcaption></figure></td>\n",
       "</tr></table>\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "import (\n",
    "    \"fmt\"\n",
    "    \"strings\"\n",
    "\t\"strconv\"\n",
    "    \"github.com/gomlx/gopjrt/dtypes\"\n",
    "    \"github.com/gomlx/gomlx/backends\"\n",
    "    \"github.com/gomlx/gomlx/examples/mnist\"\n",
    "    \"github.com/gomlx/gomlx/pkg/core/shapes\"\n",
    "    \"github.com/gomlx/gomlx/pkg/core/tensors/images\"\n",
    "    \"github.com/janpfeifer/gonb/gonbui\"\n",
    "\n",
    "    _ \"github.com/gomlx/gomlx/backends/default\"\n",
    ")\n",
    "\n",
    "var (\n",
    "    // Model DType, used everywhere.\n",
    "    DType = dtypes.Float32\n",
    ")\n",
    "\n",
    "// sampleToNotebook generates a sample of MNIST in a GoNB Jupyter Notebook.\n",
    "func sampleToNotebook() {\n",
    "    // Load data into tensors.\n",
    "    backend := backends.MustNew()\n",
    "    if ds, err := mnist.NewDataset(backend, \"Samples MNIST\", *flagDataDir, \"train\", DType); err != nil {\n",
    "        fmt.Printf(\"mnist.NewDataset: %v\", err)\n",
    "    } else {\n",
    "        ds.Shuffle()\n",
    "        sampleImages(ds, 10)\n",
    "    }\n",
    "   \n",
    "}\n",
    "\n",
    "// sampleTable generates and outputs one html table of samples, sampling rows x cols from the images/labels provided.\n",
    "func sampleImages(ds train.Dataset, numImages int) {\n",
    "    gonbui.DisplayHTML(fmt.Sprintf(\"<p>%s</p>\\n\", ds.Name()))\n",
    "    \n",
    "    parts := make([]string, 0, numImages+5) // Leave last part empty.\n",
    "    parts = append(parts, \"<table><tr>\")\n",
    "    for ii := 0; ii < numImages; ii++ {\n",
    "        _, inputs, labels := must.M3(ds.Yield())\n",
    "        imgTensor := inputs[0]\n",
    "        img := images.ToImage().Single(imgTensor)\n",
    "        label := labels[0].Value().([]int8)\n",
    "    \n",
    "        imgSrc := must.M1(gonbui.EmbedImageAsPNGSrc(img))\n",
    "        size := imgTensor.Shape().Dimensions[0]\n",
    "        parts = append(\n",
    "            parts, \n",
    "            fmt.Sprintf(`<td><figure style=\"padding:4px;text-align: center;\"><img width=\"%d\" height=\"%d\" src=\"%s\">` + \n",
    "                        `<figcaption style=\"text-align: center;\">(%d)</figcaption></figure></td>`, \n",
    "                        size*2, size*2, imgSrc, label),\n",
    "        )\n",
    "    }\n",
    "    parts = append(parts, \"</tr></table>\", \"\")\n",
    "    gonbui.DisplayHTML(strings.Join(parts, \"\\n\"))\n",
    "}\n",
    "\n",
    "%%\n",
    "AssertDownloaded()\n",
    "sampleToNotebook()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Training on MNIST\n",
    "\n",
    "### Models Support\n",
    "\n",
    "1. `flagModel` defines the model type, out of `validModels` options.\n",
    "1. `createDefaultContext` creates a context and set the default values for the MNIST models. \n",
    "1. `contextFromSettings` uses `createDefaultContext` and incorporate changes passed by the `-set` flag."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Models: [\"linear\" \"cnn\"]\n",
      "Parameters set (-set): [\"batch_size\" \"model\" \"train_steps\"]\n",
      "\t\"/activation\": (string) relu\n",
      "\t\"/adam_dtype\": (string) \n",
      "\t\"/adam_epsilon\": (float64) 1e-07\n",
      "\t\"/batch_size\": (int) 17\n",
      "\t\"/cnn_dropout_rate\": (float64) 0.5\n",
      "\t\"/cnn_normalization\": (string) layer\n",
      "\t\"/cosine_schedule_steps\": (int) 0\n",
      "\t\"/dropout_rate\": (float64) 0.5\n",
      "\t\"/eval_batch_size\": (int) 1000\n",
      "\t\"/l1_regularization\": (float64) 0\n",
      "\t\"/l2_regularization\": (float64) 0\n",
      "\t\"/learning_rate\": (float64) 0.0001\n",
      "\t\"/loss\": (string) sparse_cross_logits\n",
      "\t\"/model\": (string) 'cnn'\n",
      "\t\"/nan_logger\": (bool) false\n",
      "\t\"/num_checkpoints\": (int) 3\n",
      "\t\"/optimizer\": (string) adamw\n",
      "\t\"/plots\": (bool) false\n",
      "\t\"/train_steps\": (int) 10\n",
      "\t\"/triplet_loss_margin\": (float64) 0.5\n",
      "\t\"/triplet_loss_mining_strategy\": (string) Hard\n",
      "\t\"/triplet_loss_pairwise_distance_metric\": (string) L2\n"
     ]
    }
   ],
   "source": [
    "import (\n",
    "    \"flags\"\n",
    "    \n",
    "    \"github.com/gomlx/gomlx/pkg/ml/layers\"\n",
    "    \"github.com/gomlx/gomlx/ui/commandline\"\n",
    "    \"github.com/gomlx/gomlx/pkg/ml/train/optimizers\"\n",
    "    \"github.com/gomlx/gomlx/examples/mnist\"\n",
    "    \"github.com/gomlx/gomlx/pkg/ml/context\"\n",
    ")\n",
    "\n",
    "var (\n",
    "\tflagEval      = flag.Bool(\"eval\", true, \"Whether to evaluate the model on the validation data in the end.\")\n",
    ")\n",
    "\n",
    "// settings is bound to a \"-set\" flag to be used to set context hyperparameters.\n",
    "var settings = commandline.CreateContextSettingsFlag(CreateDefaultContext(), \"set\")\n",
    "\n",
    "// createDefaultContext sets the context with default hyperparameters\n",
    "func CreateDefaultContext() *context.Context {\n",
    "\tctx := context.New()\n",
    "\tctx.RngStateReset()\n",
    "\tctx.SetParams(map[string]any{\n",
    "\t\t// Model type to use\n",
    "\t\t\"model\":           \"linear\",\n",
    "\t\t\"loss\":            \"sparse_cross_logits\",\n",
    "\t\t\"num_checkpoints\": 3,\n",
    "\t\t\"train_steps\":     4000,\n",
    "\n",
    "\t\t// batch_size for training.\n",
    "\t\t\"batch_size\": 600,\n",
    "\n",
    "\t\t// eval_batch_size can be larger than training, it's more efficient.\n",
    "\t\t\"eval_batch_size\": 1000,\n",
    "\n",
    "\t\t// Debug parameters.\n",
    "\t\t\"nan_logger\": false, // Trigger nan error as soon as it happens -- expensive, but helps debugging.\n",
    "\n",
    "\t\t// \"plots\" trigger generating intermediary eval data for plotting, and if running in GoNB, to actually\n",
    "\t\t// draw the plot with Plotly.\n",
    "\t\t//\n",
    "\t\t// From the command-line, an easy way to monitor the metrics being generated during the training of a model\n",
    "\t\t// is using the gomlx_checkpoints tool:\n",
    "\t\t//\n",
    "\t\t//\t$ gomlx_checkpoints --metrics --metrics_labels --metrics_types=accuracy  --metrics_names='E(Tra)/#loss,E(Val)/#loss' --loop=3s \"<checkpoint_path>\"\n",
    "\t\tplotly.ParamPlots: false,\n",
    "\n",
    "\t\toptimizers.ParamOptimizer:       \"adamw\",\n",
    "\t\toptimizers.ParamLearningRate:    1e-4,\n",
    "\t\toptimizers.ParamAdamEpsilon:     1e-7,\n",
    "\t\toptimizers.ParamAdamDType:       \"\",\n",
    "\t\tcosineschedule.ParamPeriodSteps: 0,\n",
    "\t\tactivations.ParamActivation:     \"relu\",\n",
    "\t\tlayers.ParamDropoutRate:         0.5,\n",
    "\t\tregularizers.ParamL2:            0.0,\n",
    "\t\tregularizers.ParamL1:            0.0,\n",
    "\n",
    "\t\t// CNN\n",
    "\t\t\"cnn_dropout_rate\":  0.5,\n",
    "\t\t\"cnn_normalization\": \"layer\", // \"layer\" or \"batch\".\n",
    "\n",
    "\t\t// Triplet\n",
    "\t\tlosses.ParamTripletLossPairwiseDistanceMetric: \"L2\",\n",
    "\t\tlosses.ParamTripletLossMiningStrategy:         \"Hard\",\n",
    "\t\tlosses.ParamTripletLossMargin:                 0.5,\n",
    "\t})\n",
    "\treturn ctx\n",
    "}\n",
    "\n",
    "// ContextFromSettings is the default context (createDefaultContext) changed by -set flag.\n",
    "func ContextFromSettings() (ctx *context.Context, paramsSet []string) {\n",
    "    ctx = mnist.CreateDefaultContext()\n",
    "    paramsSet = must.M1(commandline.ParseContextSettings(ctx, *settings))\n",
    "    return\n",
    "}\n",
    "\n",
    "// Let's test that we can set hyperparameters by setting it in the \"-set\" flag:\n",
    "%% -set=\"batch_size=17;model='cnn';train_steps=10\"\n",
    "fmt.Printf(\"Models: %q\\n\", mnist.ModelList)\n",
    "ctx, parametersSet := ContextFromSettings()\n",
    "fmt.Printf(\"Parameters set (-set): %q\\n\", parametersSet)\n",
    "fmt.Println(commandline.SprintContextSettings(ctx))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Linear model\n",
    "\n",
    "A linear model can easily get to ~92% accuracy (a random model would do 10%) with 4000 steps.\n",
    "\n",
    "Later we are going to define a CNN model to compare, and we just set a placeholder model here for now.\n",
    "\n",
    "> **Note**: \n",
    ">\n",
    "> * The code is here just to exemplify. We are actually using the same code from the [`mnist`](https://github.com/gomlx/gomlx/tree/main/examples/mnist) package."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Logits shape for batch_size=10: (Float32)[10 10]\n"
     ]
    }
   ],
   "source": [
    "import (\n",
    "\t\"github.com/gomlx/gomlx/backends\"\n",
    "\t. \"github.com/gomlx/gomlx/pkg/core/graph\"\n",
    "\t\"github.com/gomlx/gomlx/pkg/ml/context\"\n",
    "\t\"github.com/gomlx/gomlx/pkg/ml/layers\"\n",
    ")\n",
    "\n",
    "var _ = NewGraph  // Make sure the graph package is in use.\n",
    "\n",
    "// LinearModelGraph builds a simple  model logistic model\n",
    "// It returns the logit, not the predictions, which works with most losses with shape `[batch_size, NumClasses]`.\n",
    "// inputs: only one tensor, with shape `[batch_size, width, height, depth]`.\n",
    "func LinearModelGraph(ctx *context.Context, spec any, inputs []*Node) []*Node {\n",
    "\tctx = ctx.In(\"model\") // Create the model by default under the \"/model\" scope.\n",
    "\tbatchSize := inputs[0].Shape().Dimensions[0]\n",
    "\tembeddings := Reshape(inputs[0], batchSize, -1)\n",
    "\tlogits := layers.DenseWithBias(ctx, embeddings, mnist.NumClasses)\n",
    "\treturn []*Node{logits}\n",
    "}\n",
    "\n",
    "%% -set=\"batch_size=10\"\n",
    "// Let's test that the logits are coming out with the right shape: we want [batch_size, 10], since there are 10 classes.\n",
    "AssertDownloaded()\n",
    "ctx, _ := ContextFromSettings()\n",
    "g := NewGraph(backends.MustNew(), \"placeholder\")\n",
    "batchSize := context.GetParamOr(ctx, \"batch_size\", int(100))\n",
    "logits := LinearModelGraph(ctx, nil, []*Node{Parameter(g, \"images\", shapes.Make(DType, batchSize, mnist.Height, mnist.Width, mnist.Depth))})\n",
    "fmt.Printf(\"Logits shape for batch_size=%d: %s\\n\", batchSize, logits[0].Shape())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Training Loop\n",
    "\n",
    "With a model function defined, we use the training loop create for the MNIST.\n",
    "\n",
    "The trainer is provided in the [`mnist` package](https://github.com/gomlx/gomlx/tree/main/examples/mnist). It is straight forward (and almost the same for every different project) and does the following for us:\n",
    "\n",
    "- If a checkpoing is given (--checkpoint) and it has previously saved model, it loads hyperparmeters and trained variables.\n",
    "- Create trainer: with selected model function (see [Linear model](#Linear-model) and [Linear model for MNIST](#CNN-model-for-MNIST) sections), optimizer, loss and metrics.\n",
    "- Create a `train.Loop` and attach to it a progressbar, a periodic checkpoint saver and a plotter (`--set=\"plots=true\"`).\n",
    "- Train the selected number of train steps.\n",
    "- Report results.\n",
    "\n",
    "Below we train 4000 steps with the default settings just to check things are working."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Training linear model:\n",
      "Backend stablehlo: stablehlo:cuda - PJRT \"cuda\" plugin (/home/janpf/.local/lib/gomlx/pjrt/pjrt_c_api_cuda_plugin.so) v0.76 [StableHLO]\n",
      "\t- checkpoint in /home/janpf/work/mnist/linear\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<div id=\"d017db33\"></div>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": []
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\t- restarting from global step 4000\n",
      "\t - target train_steps=4000 already reached. To train further, set a number additional to current global step.\n",
      "\n",
      "Results on train:\n",
      "\tMean Loss+Regularization (#loss+): 0.287\n",
      "\tMean Loss (#loss): 0.287\n",
      "\tMean Accuracy (#acc): 92.07%\n",
      "Results on test:\n",
      "\tMean Loss+Regularization (#loss+): 0.284\n",
      "\tMean Loss (#loss): 0.284\n",
      "\tMean Accuracy (#acc): 92.09%\n",
      "\n"
     ]
    }
   ],
   "source": [
    "var flagCheckpoint = flag.String(\"checkpoint\", \"\", \"Directory save and load checkpoints from. If left empty, no checkpoints are created.\")\n",
    "\n",
    "// trainModel with hyperparameters configured with `-set=...`.\n",
    "func trainModel() {\n",
    "    ctx, paramsSet := ContextFromSettings()\n",
    "    must.M(mnist.TrainModel(ctx, *flagDataDir, *flagCheckpoint, paramsSet))\n",
    "}\n",
    "\n",
    "// Train 50 steps, only to test things are working. No plots.\n",
    "%% --checkpoint=linear  --set=\"model=linear;train_steps=4000;plots=true\"\n",
    "trainModel()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### CNN Model for MNIST\n",
    "\n",
    "Let's now properly define a CNN model to compare.\n",
    "\n",
    "The model was built following a [Deep MNIST for Experts](https://chromium.googlesource.com/external/github.com/tensorflow/tensorflow/+/r0.10/tensorflow/g3doc/tutorials/mnist/pros/index.md)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Logits shape for batch_size=10: (Float32)[10 10]\n"
     ]
    }
   ],
   "source": [
    "// CnnModelGraph builds the CNN model for our demo.\n",
    "// It returns the logit, not the predictions, which works with most losses with shape `[batch_size, NumClasses]`.\n",
    "// inputs: only one tensor, with shape `[batch_size, width, height, depth]`.\n",
    "func CnnModelGraph(ctx *context.Context, spec any, inputs []*Node) []*Node {\n",
    "\tctx = ctx.In(\"model\") // Create the model by default under the \"/model\" scope.\n",
    "\tembeddings := CnnEmbeddings(ctx, inputs[0])\n",
    "\tlogits := layers.Dense(ctx, embeddings, true, mnist.NumClasses)\n",
    "\treturn []*Node{logits}\n",
    "}\n",
    "\n",
    "func CnnEmbeddings(ctx *context.Context, images *Node) *Node {\n",
    "\tbatchSize := images.Shape().Dimensions[0]\n",
    "\tg := images.Graph()\n",
    "\tdtype := images.DType()\n",
    "\n",
    "\tlayerIdx := 0\n",
    "\tnextCtx := func(name string) *context.Context {\n",
    "\t\tnewCtx := ctx.Inf(\"%03d_%s\", layerIdx, name)\n",
    "\t\tlayerIdx++\n",
    "\t\treturn newCtx\n",
    "\t}\n",
    "\t// Dropout.\n",
    "\tdropoutRate := context.GetParamOr(ctx, \"cnn_dropout_rate\", -1.0)\n",
    "\tif dropoutRate < 0 {\n",
    "\t\tdropoutRate = context.GetParamOr(ctx, layers.ParamDropoutRate, 0.0)\n",
    "\t}\n",
    "\tvar dropoutNode *Node\n",
    "\tif dropoutRate > 0.0 {\n",
    "\t\tdropoutNode = Scalar(g, dtype, dropoutRate)\n",
    "\t}\n",
    "\n",
    "\timages = layers.Convolution(nextCtx(\"conv\"), images).Filters(32).KernelSize(3).PadSame().Done()\n",
    "\timages.AssertDims(batchSize, 28, 28, 32)\n",
    "\timages = activations.Relu(images)\n",
    "\timages = normalizeCNN(nextCtx(\"norm\"), images)\n",
    "\timages = MaxPool(images).Window(2).Done()\n",
    "\timages.AssertDims(batchSize, 14, 14, 32)\n",
    "\n",
    "\timages = layers.Convolution(nextCtx(\"conv\"), images).Filters(64).KernelSize(3).PadSame().Done()\n",
    "\timages.AssertDims(batchSize, 14, 14, 64)\n",
    "\timages = activations.Relu(images)\n",
    "\timages = normalizeCNN(nextCtx(\"norm\"), images)\n",
    "\timages = MaxPool(images).Window(2).Done()\n",
    "\timages = layers.DropoutNormalize(nextCtx(\"dropout\"), images, dropoutNode, true)\n",
    "\timages.AssertDims(batchSize, 7, 7, 64)\n",
    "\n",
    "\t// Flatten images\n",
    "\timages = Reshape(images, batchSize, -1)\n",
    "\treturn images\n",
    "}\n",
    "\n",
    "func normalizeCNN(ctx *context.Context, logits *Node) *Node {\n",
    "\tnormalizationType := context.GetParamOr(ctx, \"cnn_normalization\", \"none\")\n",
    "\tswitch normalizationType {\n",
    "\tcase \"layer\":\n",
    "\t\tif logits.Rank() == 2 {\n",
    "\t\t\treturn layers.LayerNormalization(ctx, logits, -1).Done()\n",
    "\t\t} else if logits.Rank() == 4 {\n",
    "\t\t\treturn layers.LayerNormalization(ctx, logits, 2, 3).Done()\n",
    "\t\t} else {\n",
    "\t\t\treturn logits\n",
    "\t\t}\n",
    "\tcase \"batch\":\n",
    "\t\treturn batchnorm.New(ctx, logits, -1).Done()\n",
    "\tcase \"none\", \"\":\n",
    "\t\treturn logits\n",
    "\tdefault:\n",
    "\t\texceptions.Panicf(\"invalid normalization type %q -- set it with parameter %q\", normalizationType, \"cnn_normalization\")\n",
    "\t\treturn nil\n",
    "\t}\n",
    "}\n",
    "%% -set=\"batch_size=10\"\n",
    "// Let's test that the logits are coming out with the right shape: we want [batch_size, 10], since there are 10 classes.\n",
    "AssertDownloaded()\n",
    "ctx, _ := ContextFromSettings()\n",
    "g := NewGraph(backends.MustNew(), \"placeholder\")\n",
    "batchSize := context.GetParamOr(ctx, \"batch_size\", int(100))\n",
    "logits := CnnModelGraph(ctx, nil, []*Node{Parameter(g, \"images\", shapes.Make(DType, batchSize, mnist.Height, mnist.Width, mnist.Depth))})\n",
    "fmt.Printf(\"Logits shape for batch_size=%d: %s\\n\", batchSize, logits[0].Shape())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### CNN Model Training\n",
    "\n",
    "Let's train the CNN for real this time. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "// Remove a previously trained model\n",
    "!rm -rf ~/work/mnist/cnn  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Training cnn model:\n",
      "Backend stablehlo: stablehlo:cuda - PJRT \"cuda\" plugin (/home/janpf/.local/lib/gomlx/pjrt/pjrt_c_api_cuda_plugin.so) v0.76 [StableHLO]\n",
      "\t- checkpoint in /home/janpf/work/mnist/cnn\n"
     ]
    },
    {
     "data": {
      "text/html": []
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "      \u001b[1m   0% [........................................] (17 steps/s) [1s:3m52s]\u001b[0m [step=35] [loss+=2.05] [~loss+=2.97] [~loss=2.97] [~acc=24.69%]        "
     ]
    },
    {
     "data": {
      "text/html": []
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "      \u001b[1m 100% [========================================] (101 steps/s)\u001b[0m [step=3999] [loss+=0.0353] [~loss+=0.0464] [~loss=0.0464] [~acc=98.52%]        ]         \n",
      "\n",
      "\t- saving checkpoint@4000\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<p><b>Metric: accuracy</b></p>\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<div id=\"05313571\"></div>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<script charset=\"UTF-8\">\n",
       "(() => {\n",
       "\tconst src=\"https://cdn.plot.ly/plotly-2.34.0.min.js\";\n",
       "\tvar runJSFn = function(module) {\n",
       "\t\t\n",
       "\tif (!module) {\n",
       "\t\tmodule = window.Plotly;\n",
       "\t}\n",
       "\tlet data = JSON.parse('{\"data\":[{\"type\":\"scatter\",\"line\":{\"shape\":\"linear\"},\"mode\":\"lines+markers\",\"name\":\"Train: Moving Average Accuracy\",\"x\":[11,23,37,54,74,98,127,162,204,254,314,386,472,575,699,848,1027,1242,1500,1810,2182,2628,3163,3805,4000],\"y\":[0.13651515543460846,0.18818844854831696,0.25090092420578003,0.3171297013759613,0.3849325180053711,0.44559526443481445,0.5095990896224976,0.5791418552398682,0.6507223844528198,0.7208177447319031,0.782677173614502,0.8312121629714966,0.8709171414375305,0.9003956913948059,0.9209686517715454,0.9366459250450134,0.9465472102165222,0.9567850828170776,0.9641646146774292,0.9690183997154236,0.9744048118591309,0.978425920009613,0.981241762638092,0.9848059415817261,0.9851862788200378]},{\"type\":\"scatter\",\"line\":{\"shape\":\"linear\"},\"mode\":\"lines+markers\",\"name\":\"Mean Accuracy on train\",\"x\":[11,23,37,54,74,98,127,162,204,254,314,386,472,575,699,848,1027,1242,1500,1810,2182,2628,3163,3805,4000],\"y\":[0.2895500063896179,0.5180833339691162,0.6686999797821045,0.7509666681289673,0.8029666543006897,0.8344333171844482,0.8623499870300293,0.8845500349998474,0.9032833576202393,0.9183499813079834,0.9311333298683167,0.9422500133514404,0.9520166516304016,0.9597166776657104,0.9660000205039978,0.971583366394043,0.9767000079154968,0.9800000190734863,0.9830666780471802,0.9854666590690613,0.9881333708763123,0.9900000095367432,0.9917666912078857,0.9932667016983032,0.9939666986465454]},{\"type\":\"scatter\",\"line\":{\"shape\":\"linear\"},\"mode\":\"lines+markers\",\"name\":\"Mean Accuracy on test\",\"x\":[11,23,37,54,74,98,127,162,204,254,314,386,472,575,699,848,1027,1242,1500,1810,2182,2628,3163,3805,4000],\"y\":[0.2979000210762024,0.5338000059127808,0.6767000555992126,0.7646000385284424,0.8109000325202942,0.8416000604629517,0.8677000403404236,0.8915000557899475,0.9097000360488892,0.9235000610351562,0.9361000657081604,0.9466000199317932,0.9551000595092773,0.9611000418663025,0.9677000641822815,0.9711000323295593,0.9745000600814819,0.9782000184059143,0.9807000756263733,0.9825000762939453,0.9860000610351562,0.9861000180244446,0.9866000413894653,0.9873000383377075,0.9875000715255737]}],\"layout\":{\"legend\":{},\"title\":{\"text\":\"accuracy\"},\"xaxis\":{\"showgrid\":true,\"title\":{\"text\":\"Steps\"},\"type\":\"log\"},\"yaxis\":{\"showgrid\":true,\"type\":\"log\"}}}');\n",
       "\tmodule.newPlot('05313571', data);\n",
       "\n",
       "\t}\n",
       "\t\n",
       "    if (typeof requirejs === \"function\") {\n",
       "        // Use RequireJS to load module.\n",
       "\t\tlet srcWithoutExtension = src.substring(0, src.lastIndexOf(\".js\"));\n",
       "        requirejs.config({\n",
       "            paths: {\n",
       "                'plotly': srcWithoutExtension\n",
       "            }\n",
       "        });\n",
       "        require(['plotly'], function(plotly) {\n",
       "            runJSFn(plotly)\n",
       "        });\n",
       "        return\n",
       "    }\n",
       "\n",
       "\tvar currentScripts = document.head.getElementsByTagName(\"script\");\n",
       "\tfor (const idx in currentScripts) {\n",
       "\t\tlet script = currentScripts[idx];\n",
       "\t\tif (script.src == src) {\n",
       "\t\t\trunJSFn(null);\n",
       "\t\t\treturn;\n",
       "\t\t}\n",
       "\t}\n",
       "\n",
       "\tvar script = document.createElement(\"script\");\n",
       "\n",
       "\tscript.charset = \"utf-8\";\n",
       "\t\n",
       "\tscript.src = src;\n",
       "\tscript.onload = script.onreadystatechange = function () { runJSFn(null); };\n",
       "\tdocument.head.appendChild(script);\t\n",
       "})();\n",
       "</script>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<p><b>Metric: loss</b></p>\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<div id=\"03291472\"></div>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<script charset=\"UTF-8\">\n",
       "(() => {\n",
       "\tconst src=\"https://cdn.plot.ly/plotly-2.34.0.min.js\";\n",
       "\tvar runJSFn = function(module) {\n",
       "\t\t\n",
       "\tif (!module) {\n",
       "\t\tmodule = window.Plotly;\n",
       "\t}\n",
       "\tlet data = JSON.parse('{\"data\":[{\"type\":\"scatter\",\"line\":{\"shape\":\"linear\"},\"mode\":\"lines+markers\",\"name\":\"Train: Batch Loss+Regularization\",\"x\":[11,23,37,54,74,98,127,162,204,254,314,386,472,575,699,848,1027,1242,1500,1810,2182,2628,3163,3805,4000],\"y\":[3.437448024749756,2.5676777362823486,2.0794734954833984,1.5502299070358276,1.3108669519424438,0.9957395792007446,0.8655920624732971,0.6982430219650269,0.6172050833702087,0.5686812996864319,0.3627985119819641,0.37803152203559875,0.32124048471450806,0.20560772716999054,0.19341517984867096,0.22084788978099823,0.1669515073299408,0.12788164615631104,0.10599105805158615,0.0831967294216156,0.05175691843032837,0.09292913228273392,0.06877386569976807,0.033459197729825974,0.03534816578030586]},{\"type\":\"scatter\",\"line\":{\"shape\":\"linear\"},\"mode\":\"lines+markers\",\"name\":\"Train: Moving Average Loss+Regularization\",\"x\":[11,23,37,54,74,98,127,162,204,254,314,386,472,575,699,848,1027,1242,1500,1810,2182,2628,3163,3805,4000],\"y\":[3.9154558181762695,3.3821868896484375,2.9506025314331055,2.592423677444458,2.275583505630493,2.00748872756958,1.7425988912582397,1.4672749042510986,1.1971473693847656,0.9427398443222046,0.7229757905006409,0.559377133846283,0.42746502161026,0.3291900157928467,0.2597971260547638,0.20746728777885437,0.17495931684970856,0.1398036777973175,0.1163683831691742,0.09953074902296066,0.08219995349645615,0.06823853403329849,0.059023793786764145,0.046948954463005066,0.0463678278028965]},{\"type\":\"scatter\",\"line\":{\"shape\":\"linear\"},\"mode\":\"lines+markers\",\"name\":\"Train: Moving Average Loss\",\"x\":[11,23,37,54,74,98,127,162,204,254,314,386,472,575,699,848,1027,1242,1500,1810,2182,2628,3163,3805,4000],\"y\":[3.9154558181762695,3.3821868896484375,2.9506025314331055,2.592423677444458,2.275583505630493,2.00748872756958,1.7425988912582397,1.4672749042510986,1.1971473693847656,0.9427398443222046,0.7229757905006409,0.559377133846283,0.42746502161026,0.3291900157928467,0.2597971260547638,0.20746728777885437,0.17495931684970856,0.1398036777973175,0.1163683831691742,0.09953074902296066,0.08219995349645615,0.06823853403329849,0.059023793786764145,0.046948954463005066,0.0463678278028965]},{\"type\":\"scatter\",\"line\":{\"shape\":\"linear\"},\"mode\":\"lines+markers\",\"name\":\"Mean Loss+Regularization on train\",\"x\":[11,23,37,54,74,98,127,162,204,254,314,386,472,575,699,848,1027,1242,1500,1810,2182,2628,3163,3805,4000],\"y\":[2.0886573791503906,1.4657198190689087,1.040674090385437,0.7802774310112,0.615855872631073,0.5095051527023315,0.4253307282924652,0.3569593131542206,0.3016677498817444,0.25815969705581665,0.2194766253232956,0.18421189486980438,0.15435241162776947,0.12969444692134857,0.10937698930501938,0.09191607683897018,0.07744959741830826,0.06621090322732925,0.05587415397167206,0.047532569617033005,0.03976575657725334,0.03354785591363907,0.027669360861182213,0.02188281901180744,0.02014850452542305]},{\"type\":\"scatter\",\"line\":{\"shape\":\"linear\"},\"mode\":\"lines+markers\",\"name\":\"Mean Loss on train\",\"x\":[11,23,37,54,74,98,127,162,204,254,314,386,472,575,699,848,1027,1242,1500,1810,2182,2628,3163,3805,4000],\"y\":[2.0886573791503906,1.4657198190689087,1.040674090385437,0.7802774310112,0.615855872631073,0.5095051527023315,0.4253307282924652,0.3569593131542206,0.3016677498817444,0.25815969705581665,0.2194766253232956,0.18421189486980438,0.15435241162776947,0.12969444692134857,0.10937698930501938,0.09191607683897018,0.07744959741830826,0.06621090322732925,0.05587415397167206,0.047532569617033005,0.03976575657725334,0.03354785591363907,0.027669360861182213,0.02188281901180744,0.02014850452542305]},{\"type\":\"scatter\",\"line\":{\"shape\":\"linear\"},\"mode\":\"lines+markers\",\"name\":\"Mean Loss+Regularization on test\",\"x\":[11,23,37,54,74,98,127,162,204,254,314,386,472,575,699,848,1027,1242,1500,1810,2182,2628,3163,3805,4000],\"y\":[2.0496866703033447,1.4340676069259644,1.0187486410140991,0.7604366540908813,0.5991008877754211,0.4929656386375427,0.41016945242881775,0.34366172552108765,0.28764936327934265,0.24493952095508575,0.2085978388786316,0.17696037888526917,0.1478835940361023,0.12506549060344696,0.10692352801561356,0.09157262742519379,0.07919479161500931,0.0685657262802124,0.0596090592443943,0.05271182581782341,0.04599021375179291,0.04332658648490906,0.039714131504297256,0.03546500951051712,0.03558168560266495]},{\"type\":\"scatter\",\"line\":{\"shape\":\"linear\"},\"mode\":\"lines+markers\",\"name\":\"Mean Loss on test\",\"x\":[11,23,37,54,74,98,127,162,204,254,314,386,472,575,699,848,1027,1242,1500,1810,2182,2628,3163,3805,4000],\"y\":[2.0496866703033447,1.4340676069259644,1.0187486410140991,0.7604366540908813,0.5991008877754211,0.4929656386375427,0.41016945242881775,0.34366172552108765,0.28764936327934265,0.24493952095508575,0.2085978388786316,0.17696037888526917,0.1478835940361023,0.12506549060344696,0.10692352801561356,0.09157262742519379,0.07919479161500931,0.0685657262802124,0.0596090592443943,0.05271182581782341,0.04599021375179291,0.04332658648490906,0.039714131504297256,0.03546500951051712,0.03558168560266495]}],\"layout\":{\"legend\":{},\"title\":{\"text\":\"loss\"},\"xaxis\":{\"showgrid\":true,\"title\":{\"text\":\"Steps\"},\"type\":\"log\"},\"yaxis\":{\"showgrid\":true,\"type\":\"log\"}}}');\n",
       "\tmodule.newPlot('03291472', data);\n",
       "\n",
       "\t}\n",
       "\t\n",
       "    if (typeof requirejs === \"function\") {\n",
       "        // Use RequireJS to load module.\n",
       "\t\tlet srcWithoutExtension = src.substring(0, src.lastIndexOf(\".js\"));\n",
       "        requirejs.config({\n",
       "            paths: {\n",
       "                'plotly': srcWithoutExtension\n",
       "            }\n",
       "        });\n",
       "        require(['plotly'], function(plotly) {\n",
       "            runJSFn(plotly)\n",
       "        });\n",
       "        return\n",
       "    }\n",
       "\n",
       "\tvar currentScripts = document.head.getElementsByTagName(\"script\");\n",
       "\tfor (const idx in currentScripts) {\n",
       "\t\tlet script = currentScripts[idx];\n",
       "\t\tif (script.src == src) {\n",
       "\t\t\trunJSFn(null);\n",
       "\t\t\treturn;\n",
       "\t\t}\n",
       "\t}\n",
       "\n",
       "\tvar script = document.createElement(\"script\");\n",
       "\n",
       "\tscript.charset = \"utf-8\";\n",
       "\t\n",
       "\tscript.src = src;\n",
       "\tscript.onload = script.onreadystatechange = function () { runJSFn(null); };\n",
       "\tdocument.head.appendChild(script);\t\n",
       "})();\n",
       "</script>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\t- trained to step 4000, median train step: 3965 microseconds\n",
      "\n",
      "Results on train:\n",
      "\tMean Loss+Regularization (#loss+): 0.0201\n",
      "\tMean Loss (#loss): 0.0201\n",
      "\tMean Accuracy (#acc): 99.40%\n",
      "Results on test:\n",
      "\tMean Loss+Regularization (#loss+): 0.0356\n",
      "\tMean Loss (#loss): 0.0356\n",
      "\tMean Accuracy (#acc): 98.75%\n",
      "\n"
     ]
    }
   ],
   "source": [
    "%% --checkpoint=cnn --set=\"model=cnn;train_steps=4000;plots=true\"\n",
    "trainModel()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Go (gonb)",
   "language": "go",
   "name": "gonb"
  },
  "language_info": {
   "codemirror_mode": "",
   "file_extension": ".go",
   "mimetype": "text/x-go",
   "name": "go",
   "nbconvert_exporter": "",
   "pygments_lexer": "",
   "version": "go1.25.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
